
#pragma once

#include <cassert>
#include <cmath>

// Assume Core2 compatible CPU
#include <tmmintrin.h>

//////////////////////////////////////////////////////////////////////////
// Matrices, vectors
// column major storage (same as in default glsl)
// no classy tricks here, this should be lean and mean (in order to work as a buffer)

class vec2
{
public:
	vec2();
	vec2(float _x, float _y);


	// data members
	union
	{
		struct  
		{
			float x;
			float y;
		};

		struct  
		{
			float u;
			float v;
		};
	};
};

class vec3
{
public:
    vec3();
    vec3(float x, float y, float z);
    vec3 operator+(const vec3& rhs) const;
    vec3 operator-() const;
    vec3 operator-(const vec3& rhs) const;
    vec3 operator*(float rhs) const;
    vec3& operator *=(float rhs);
    vec3& Normalize();
    float Norm();
    float x, y, z;
};

float dot(const vec3& lhs, const vec3& rhs);
vec3 cross(const vec3& lhs, const vec3& rhs);

class vec4 {
    __m128 m_elems;

    friend class mat4;
public:
    
    vec4() { m_elems=_mm_setzero_ps(); }
    vec4(const vec4& src) { m_elems=src.m_elems; }
    vec4& operator=(const vec4& src) { m_elems=src.m_elems; return *this; }
    vec4& operator=(const __m128& src) { m_elems=src; return *this; }
    float& operator[](int idx) { assert(idx<4); return m_elems.m128_f32[idx]; }
    operator float*() { return (float*)(m_elems.m128_f32); }
    operator const __m128&() const { return m_elems; }
};

class mat4 {
    __m128 m_cols[4];

public:
    mat4()
    {
        m_cols[0]=_mm_setr_ps(1,0,0,0);
        m_cols[1]=_mm_setr_ps(0,1,0,0);
        m_cols[2]=_mm_setr_ps(0,0,1,0);
        m_cols[3]=_mm_setr_ps(0,0,0,1);
    }
    mat4(float* elems) { for(int i=0;i<4;i++) m_cols[i]=_mm_loadu_ps(elems+4*i); }
    mat4(const mat4& src) { for(int i=0;i<4;i++) m_cols[i]=src.m_cols[i]; }
    mat4& operator=(const mat4& src) { for(int i=0;i<4;i++) m_cols[i]=src.m_cols[i]; return *this; }
    mat4 operator*(const mat4& m) const { mat4 tmp(m); return MulRight(tmp); }
    vec4& MulRight(vec4& v) const
    {
        __m128 tmp[4];

        tmp[0]=v.m_elems;
        tmp[1]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(1,1,1,1));
        tmp[2]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(2,2,2,2));
        tmp[3]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(3,3,3,3));
        tmp[0]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(0,0,0,0));

        tmp[0]=_mm_mul_ps(tmp[0], m_cols[0]);
        tmp[1]=_mm_mul_ps(tmp[1], m_cols[1]);
        tmp[2]=_mm_mul_ps(tmp[2], m_cols[2]);
        tmp[3]=_mm_mul_ps(tmp[3], m_cols[3]);

        tmp[0]=_mm_add_ps(tmp[0], tmp[1]);
        tmp[2]=_mm_add_ps(tmp[2], tmp[3]);
        v.m_elems=_mm_add_ps(tmp[0], tmp[2]);
        return v;
    }
    mat4& MulRight(mat4& m) const
    {
        for(int i=0;i<4;i++) {
            __m128 tmp[4];
            __m128 right=m.m_cols[i];

            tmp[0]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(0, 0, 0, 0));
            tmp[0]=_mm_mul_ps(tmp[0], m_cols[0]);
            tmp[1]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(1, 1, 1, 1));
            tmp[1]=_mm_mul_ps(tmp[1], m_cols[1]);
            tmp[2]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(2, 2, 2, 2));
            tmp[2]=_mm_mul_ps(tmp[2], m_cols[2]);
            tmp[3]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(3, 3, 3, 3));
            tmp[3]=_mm_mul_ps(tmp[3], m_cols[3]);

            tmp[0]=_mm_add_ps(tmp[0], tmp[1]);
            tmp[2]=_mm_add_ps(tmp[2], tmp[3]);
            m.m_cols[i]=_mm_add_ps(tmp[0], tmp[2]);
        }

        return m;
    }
    mat4& operator*=(const mat4& m)
    {
        mat4 tmp(m);
        MulRight(tmp);
        m_cols[0]=tmp.m_cols[0];
        m_cols[1]=tmp.m_cols[1];
        m_cols[2]=tmp.m_cols[2];
        m_cols[3]=tmp.m_cols[3];

        return *this;
    }
    mat4& operator*=(float a)
    {
        __m128 tmp=_mm_set1_ps(a);
        m_cols[0]=_mm_mul_ps(m_cols[0], tmp);
        m_cols[1]=_mm_mul_ps(m_cols[1], tmp);
        m_cols[2]=_mm_mul_ps(m_cols[2], tmp);
        m_cols[3]=_mm_mul_ps(m_cols[3], tmp);
        return *this;
    }
    mat4& operator+=(mat4& m)
    {
        m_cols[0]=_mm_add_ps(m_cols[0], m.m_cols[0]);
        m_cols[1]=_mm_add_ps(m_cols[1], m.m_cols[1]);
        m_cols[2]=_mm_add_ps(m_cols[2], m.m_cols[2]);
        m_cols[3]=_mm_add_ps(m_cols[3], m.m_cols[3]);
        return *this;
    }
    mat4& operator-=(mat4& m)
    {
        m_cols[0]=_mm_sub_ps(m_cols[0], m.m_cols[0]);
        m_cols[1]=_mm_sub_ps(m_cols[1], m.m_cols[1]);
        m_cols[2]=_mm_sub_ps(m_cols[2], m.m_cols[2]);
        m_cols[3]=_mm_sub_ps(m_cols[3], m.m_cols[3]);
        return *this;
    }
    float* NormalMatrix(float* dst) const   // Adjugate adjoint 4x3
    {
        __m128 u=m_cols[1];
        __m128 v=m_cols[2];

        __m128 res=_mm_mul_ps(_mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 1, 0, 2)));
        __m128 res1=_mm_mul_ps(_mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 1, 0, 2)));
        *(__m128*)(dst) = _mm_sub_ps(res, res1);

        u=v;
        v=m_cols[0];
        res=_mm_mul_ps(_mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 1, 0, 2)));
        res1=_mm_mul_ps(_mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 1, 0, 2)));
        *(__m128*)(dst+4) = _mm_sub_ps(res, res1);

        u=v;
        v=m_cols[1];
        res=_mm_mul_ps(_mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 1, 0, 2)));
        res1=_mm_mul_ps(_mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 1, 0, 2)));
        *(__m128*)(dst+8) = _mm_sub_ps(res, res1);

        return dst;
    }
    operator float*() { return (float*)m_cols; }
    float& operator()(int i, int j) { return m_cols[j].m128_f32[i]; }
    __m128& operator()(int j) { return m_cols[j]; }

    mat4 t() const;
};

inline mat4 RotationMatrix(const vec4& axis, float angle)
{
    // use Rodriguez's rotation formula
    mat4 kx;
    __m128 k;
    __m128i tmp, signs;

    k = (__m128)axis;

    signs = _mm_cvtsi32_si128(1<<31);
    tmp = _mm_castps_si128(k);

    kx(1) = _mm_castsi128_ps(_mm_xor_si128(_mm_shuffle_epi32(tmp, _MM_SHUFFLE(3, 0, 3, 2)), signs));

    signs = _mm_slli_si128(signs, 4);
    kx(2) = _mm_castsi128_ps(_mm_xor_si128(_mm_shuffle_epi32(tmp, _MM_SHUFFLE(3, 3, 0, 1)), signs));

    signs = _mm_slli_si128(signs, 4);
    kx(0) = _mm_castsi128_ps(_mm_xor_si128(_mm_shuffle_epi32(tmp, _MM_SHUFFLE(3, 1, 2, 3)), signs));

    kx(3) = _mm_setzero_ps();

    kx*=sin(angle);

    mat4 kk;
    kk(0) = _mm_mul_ps(k, _mm_shuffle_ps(k, k, _MM_SHUFFLE(3,0,0,0)));
    kk(1) = _mm_mul_ps(k, _mm_shuffle_ps(k, k, _MM_SHUFFLE(3,1,1,1)));
    kk(2) = _mm_mul_ps(k, _mm_shuffle_ps(k, k, _MM_SHUFFLE(3,2,2,2)));
    kk-=mat4();
    kk*=1-cos(angle);
    kk+=mat4();

    kk+=kx;

    return kk;
}

mat4 MatrixScale(float x, float y, float z);
mat4 MatrixLookAt(vec3& pos, vec3& traget, vec3& up);
mat4 MatrixPerspectiveFOV(float FOV, float Width, float Height, float zNear, float zFar);
mat4 TranslateMatrix( float x, float y, float z);